package com.xiashitech.agent.instrumentation.db.statement;

import com.xiashitech.agent.config.StartConfig;
import com.xiashitech.agent.instrumentation.base.AbstractAttributesExtractor;
import com.xiashitech.agent.instrumentation.utils.BaggageUtil;
import com.xiashitech.agent.instrumentation.utils.MachineAttributesUtil;
import com.xiashitech.agent.service.UserTracingService;
import com.xiashitech.agent.utils.SessionUtil;
import com.xiashitech.interfaces.agent.dto.SystemInfoDTO;
import io.opentelemetry.javaagent.shaded.instrumentation.api.instrumenter.db.SqlClientAttributesGetter;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.baggage.Baggage;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.common.AttributesBuilder;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import org.jetbrains.annotations.Nullable;

import java.util.Map;

public class XSJdbcAttributesExtractor <DbRequest, Void> extends AbstractAttributesExtractor<DbRequest, Void> {
    public String getAttributesExtractorName() {
        return XSJdbcAttributesExtractor.class.getName();
    }
    private XSJdbcAttributesExtractor(SqlClientAttributesGetter<DbRequest> getter) {
        this.getter = getter;
    }

    private final SqlClientAttributesGetter<DbRequest> getter;

    public static <DbRequest, Void> XSJdbcAttributesExtractor<DbRequest, Void> create(SqlClientAttributesGetter<DbRequest> getter) {
        return new XSJdbcAttributesExtractor(getter);
    }

    @Override
    public void onXSStart(AttributesBuilder attributes, Context context, DbRequest dbRequest) {
        Baggage baggage = SessionUtil.syncBaggage();
        baggage = BaggageUtil.mergeBaggageFromContext(baggage,context).build();
        System.out.println(Thread.currentThread().getId() + " DesenseDebug, XSJdbcAttributesExtractor onXSStart " + baggage.getEntryValue("endpoint-api"));
        attributes.put("xs.operation.db",true);
        attributes.put("xs.db.sql",((io.opentelemetry.javaagent.shaded.instrumentation.jdbc.internal.DbRequest)dbRequest).getStatement());
        attributes.put("xs.systemId", StartConfig.getParamValueFromKey(StartConfig.systemId));
        if(SystemInfoDTO.getConfig() != null) {
            attributes.put("xs.systemReg", SystemInfoDTO.getConfig().getSystemExist());
        }
        attributes.put("xs.sourceSystemId", baggage.getEntryValue("xsAuthSystemId") == null ? StartConfig.getParamValueFromKey(StartConfig.systemId):baggage.getEntryValue("xsAuthSystemId"));
        Map<String,Object> attributeMap = UserTracingService.supplementUserTracing(baggage,context,dbRequest);
        if(attributeMap != null) {
            attributeMap.forEach((attrKey,attrValue) -> {
                setAttributes(attributes,attrKey,attrValue);
            });
        }
    }

    @Override
    public void onXSEnd(AttributesBuilder attributes, Context context, DbRequest dbRequest, @Nullable Void aVoid, @Nullable Throwable throwable) {
        MachineAttributesUtil.fillMachineAttributes(attributes);
    }

    private void setAttributes(AttributesBuilder attributes, String attrKey, Object attrValue) {
        if(attrValue != null) {
            if(attrValue instanceof Boolean || "true".equalsIgnoreCase(attrValue.toString()) || "false".equalsIgnoreCase(attrValue.toString())) {
                attributes.put(attrKey,Boolean.valueOf(attrValue.toString()));
            } else {
                attributes.put(attrKey,attrValue.toString());
            }
        } else {
            attributes.put(attrKey,"");
        }
    }
}
